from robocasa.environments.kitchen.kitchen import *


class HeatMug(Kitchen):
    def __init__(self, *args, **kwargs):        
        super().__init__( *args, **kwargs)

    def _setup_kitchen_references(self):
        super()._setup_kitchen_references()
        self.microwave = self.register_fixture_ref("microwave", dict(id=FixtureType.MICROWAVE))
        self.cab = self.register_fixture_ref("cab", dict(id=FixtureType.CABINET_TOP, ref=self.microwave))
        self.init_robot_base_pos = self.cab

    def get_ep_meta(self):
        ep_meta = super().get_ep_meta()
        ep_meta["lang"] = f"pick the mug from the cabinet and place it inside the microwave. Then close the microwave"
        return ep_meta
    
    def _reset_internal(self):
        """
        Resets simulation internal configurations.
        """
        super()._reset_internal()
        self.cab.set_door_state(min=0.90, max=1.0, env=self, rng=self.rng)
        self.microwave.set_door_state(min=0.90, max=1.0, env=self, rng=self.rng)
    
    def _get_obj_cfgs(self):
        cfgs = []
        cfgs.append(dict(
            name="obj",
            obj_groups="mug",
            graspable=True,
            placement=dict(
                fixture=self.cab,
                size=(0.50, 0.20),
                pos=(0, -1.0),
            ),
        ))

        cfgs.append(dict(
            name="distr_cab",
            obj_groups="all",
            placement=dict(
                fixture=self.cab,
                size=(1.0, 0.20),
                pos=(0.0, 1.0),
                offset=(0.0, 0.0),
            ),
        ))

        return cfgs

    def _check_success(self):
        gripper_obj_far = OU.gripper_obj_far(self)
        obj_in_microwave = OU.obj_inside_of(self, "obj", self.microwave)
        door_state = self.microwave.get_door_state(self)["door"]
        door_closed = door_state <= 0.005

        return obj_in_microwave and gripper_obj_far and door_closed